import pymupdf 
import os
import re

# 找出工作目录下的所有PDF文件
def find_all_pdf():
    my_path = './/ieee_Quan'
    pdf_list = []
    for root, dirs, files in os.walk(my_path):
        for file in files:
            if file.endswith('.pdf'):
                pdf_full_path = os.path.join(root, file)
                pdf_list.append(pdf_full_path)
                # print(pdf_full_path)
    return pdf_list



def parse_pdf_references(pdf_path, reference_pattern):
    num_words = 10 # 参考文献内容只保留前num_words个词
    doc = pymupdf.open(pdf_path)
    references = []
    # 遍历PDF的每一页
    for page_num in range(len(doc)):
        page = doc.load_page(page_num)
        text = page.get_text()
        # print(text)
        # 使用正则表达式在文本中查找匹配的参考文献
        matches = reference_pattern.findall(text)
        # print(matches)
        for match in matches:
            # 提取序号、后面的内容和年份
            index, content, year = match
            # 精简content内容只保留前num_words个词
            content_words = content.split()
            truncated_content = ' '.join(content_words[:num_words])
            # 将解析出的信息添加到列表中
            references.append({
                'index': index.strip(),
                'content': truncated_content.strip(),
                'year': year.strip()
            })
    references.sort(key=lambda x: int(x['index']))
    return references

def main():
    # 设置工作目录
    # os.chdir('c://Users//dell//Desktop//dlp_hw2//ieee_Quan')
    # 假设您已经有一个PDF文件的路径
    # pdf_path = 'A BERT base model for the analysis of Electronic Health Records from diabetic patients.pdf'
    pdf_lib=find_all_pdf()
    print(pdf_lib)
    pdf_path=pdf_lib[0]
    print(pdf_path)
    # 正则表达式模式，用于匹配文献引用
    reference_pattern = re.compile(r'\[(\d+)\]\s*(.+?)\.\s*(\d{4})\.?', re.DOTALL)
    references = parse_pdf_references(pdf_path, reference_pattern)
    # 打印解析出的参考文献信息（用于调试）
    for ref in references:
        print(ref, sep='\n')

class TestParsePdfReferences():
    def __init__(self):
        self.pdf_lib=[]
        self.pdf_path=''
        pass
    def test_parse_pdf_references(self):
        self.pdf_lib=find_all_pdf()
        # print(self.pdf_lib)
        # print('')
        for i in range(1,len(self.pdf_lib)):
            self.pdf_path=self.pdf_lib[i]
            # print(self.pdf_path)
            # print('')
            # 正则表达式模式，用于匹配文献引用
            reference_pattern = re.compile(r'\[(\d+)\]\s*(.+?)\.\s*(\d{4})\.?', re.DOTALL)
            references = parse_pdf_references(self.pdf_path, reference_pattern)
            # 打印解析出的参考文献信息（用于调试）
            with open('ieee_Quan_references.txt', 'a', encoding='utf-8') as f:    
                f.write(str(self.pdf_path))
                f.write('\n')
                for ref in references:
                    # print(ref, sep='\n')    
                    f.write(str(ref))
                    f.write('\n')
                    print('done')
                f.write('\n')
                f.close()



if __name__ == "__main__":
    test1=TestParsePdfReferences()
    test1.test_parse_pdf_references()
    # main()
    # find_all_pdf()
